#!/bin/bash

# YOLO推理服务启动脚本

set -e

# 脚本目录
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_debug() {
    echo -e "${BLUE}[DEBUG]${NC} $1"
}

# 检查Python环境
check_python() {
    log_info "检查Python环境..."
    
    if ! command -v python3 &> /dev/null; then
        log_error "Python3 未安装"
        exit 1
    fi
    
    PYTHON_VERSION=$(python3 --version | cut -d' ' -f2)
    log_info "Python版本: $PYTHON_VERSION"
    
    # 检查Python版本是否 >= 3.8
    if ! python3 -c "import sys; exit(0 if sys.version_info >= (3, 8) else 1)"; then
        log_error "需要Python 3.8或更高版本"
        exit 1
    fi
}

# 创建虚拟环境
create_venv() {
    if [ ! -d "venv" ]; then
        log_info "创建虚拟环境..."
        python3 -m venv venv
    else
        log_info "虚拟环境已存在"
    fi
}

# 激活虚拟环境
activate_venv() {
    log_info "激活虚拟环境..."
    source venv/bin/activate
}

# 安装依赖
install_dependencies() {
    log_info "安装依赖包..."
    
    # 升级pip
    pip install --upgrade pip
    
    # 安装依赖
    if [ -f "requirements.txt" ]; then
        pip install -r requirements.txt
    else
        log_error "requirements.txt 文件不存在"
        exit 1
    fi
}

# 创建模型目录
create_model_dir() {
    log_info "创建模型目录..."
    mkdir -p models
}

# 检查GPU支持
check_gpu() {
    log_info "检查GPU支持..."
    
    if command -v nvidia-smi &> /dev/null; then
        log_info "检测到NVIDIA GPU:"
        nvidia-smi --query-gpu=name,memory.total --format=csv,noheader,nounits
    else
        log_warn "未检测到NVIDIA GPU，将使用CPU推理"
    fi
}

# 启动服务器
start_server() {
    log_info "启动YOLO推理服务器..."
    
    # 默认参数
    HOST=${HOST:-"0.0.0.0"}
    PORT=${PORT:-8090}
    WORKERS=${WORKERS:-1}
    MODELS_DIR=${MODELS_DIR:-"./models"}
    
    log_info "服务器配置:"
    log_info "  主机: $HOST"
    log_info "  端口: $PORT"
    log_info "  工作进程: $WORKERS"
    log_info "  模型目录: $MODELS_DIR"
    
    # 启动服务器
    python3 yolo_inference_server.py \
        --host "$HOST" \
        --port "$PORT" \
        --workers "$WORKERS" \
        --models-dir "$MODELS_DIR" \
        "$@"
}

# 主函数
main() {
    log_info "启动YOLO推理服务..."
    
    # 检查Python环境
    check_python
    
    # 创建并激活虚拟环境
    create_venv
    activate_venv
    
    # 安装依赖
    install_dependencies
    
    # 创建模型目录
    create_model_dir
    
    # 检查GPU支持
    check_gpu
    
    # 启动服务器
    start_server "$@"
}

# 帮助信息
show_help() {
    echo "YOLO推理服务启动脚本"
    echo ""
    echo "用法: $0 [选项]"
    echo ""
    echo "环境变量:"
    echo "  HOST        服务器主机地址 (默认: 0.0.0.0)"
    echo "  PORT        服务器端口 (默认: 8000)"
    echo "  WORKERS     工作进程数 (默认: 1)"
    echo "  MODELS_DIR  模型目录 (默认: ./models)"
    echo ""
    echo "选项:"
    echo "  -h, --help     显示帮助信息"
    echo "  --reload       启用自动重载（开发模式）"
    echo ""
    echo "示例:"
    echo "  $0                    # 使用默认配置启动"
    echo "  PORT=8080 $0          # 在端口8080启动"
    echo "  $0 --reload           # 开发模式启动"
}

# 处理命令行参数
case "${1:-}" in
    -h|--help)
        show_help
        exit 0
        ;;
    *)
        main "$@"
        ;;
esac
